import os

from funcs import base_train

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from torch import nn
from GSCNN import GSCNN, JointEdgeSegLoss
from augmentations import base_aug
from loaders import JIDataset
from metrics import MeanFourMetrics
from preprocess import get_preprocess, get_preprocess_2
import argparse
import torch.optim.lr_scheduler
import yaml


def parse_args():
    config_parser = argparse.ArgumentParser(description='Training Config', add_help=False)
    config_parser.add_argument('-c', '--config', default='', type=str, metavar='FILE', help='YAML config file')
    parser = argparse.ArgumentParser(description='training process')
    parser.add_argument('log', type=str)
    parser.add_argument('--work_dirs', type=str, default='../work_dirs/building')
    parser.add_argument('--epochs', type=int, default=200)
    parser.add_argument('--bs', '--batch_size', type=int, default=30)
    parser.add_argument('--vbs', '--val_batch_size', type=int, default=5)
    parser.add_argument('--amp', action='store_true', default=True)
    parser.add_argument('--swa', action='store_true', default=True)
    parser.add_argument('--data_parallel', action='store_true', default=True)
    args_config, remaining = config_parser.parse_known_args()
    if args_config.config:
        with open(args_config.config, 'r') as f:
            cfg = yaml.safe_load(f)
            parser.set_defaults(**cfg)
    args = parser.parse_args(remaining)
    args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)
    return args, args_text


def train():
    num_classes = 2
    onehot = num_classes == 1
    preprocess = get_preprocess() if onehot else get_preprocess_2()
    args, args_text = parse_args()

    # 数据集
    train_dataset = JIDataset("train", preprocess, base_aug(), onehot=onehot)
    val_dataset = JIDataset("val", preprocess, onehot=onehot)

    # 网络
    net = GSCNN(num_classes, pretrained=True)

    # 权重位置
    checkpoint_path = None
    # checkpoint_path = '/home/peter/zze/codes/work_dirs/gaofen/fusai/c10/m5_train_7207.pth'

    optimizer = torch.optim.Adam(net.parameters(), lr=2E-4)
    criterion = JointEdgeSegLoss(num_classes)
    metrics = MeanFourMetrics(num_classes)

    scheduler = None
    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=6E-6)

    # 开始训练
    base_train(args, args_text,
               train_dataset, val_dataset,
               net, checkpoint_path,
               optimizer, criterion, metrics, scheduler)


if __name__ == '__main__':
    train()
